//
// Copyright © 2025 Hardcore Engineering Inc.
//
// Licensed under the Eclipse Public License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License. You may
// obtain a copy of the License at https://www.eclipse.org/legal/epl-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//
// See the License for the specific language governing permissions and
// limitations under the License.
//

import {
  AiTokensData,
  AiTokensUsage,
  AiTranscriptData,
  AiTranscriptUsage,
  BillingDB,
  LiveKitEgressData,
  LiveKitEgressUsageData,
  LiveKitSessionData,
  LiveKitSessionsUsageData,
  LiveKitUsageData
} from '../types'
import postgres, { type Row, Sql } from 'postgres'
import { MeasureContext, type WorkspaceUuid } from '@hcengineering/core'
import { LoggedDB } from './logged'
import { RetryDB } from './retry'
import { getMigrations } from './migrations'

const BATCH_SIZE = 100

export async function createDb (ctx: MeasureContext, connectionString: string): Promise<BillingDB> {
  const sql = postgres(connectionString, {
    max: 5,
    connection: {
      application_name: 'billing'
    },
    fetch_types: false,
    prepare: false,
    types: {
      // https://jdbc.postgresql.org/documentation/publicapi/constant-values.html
      int8: {
        to: 0,
        from: [20],
        serialize: (value: string) => value.toString(),
        parse: (value: number) => Number(value)
      }
    }
  })

  const db = await PostgresDB.create(ctx, sql)
  return new LoggedDB(ctx, new RetryDB(db, { retries: 5 }))
}

class PostgresDB implements BillingDB {
  private constructor (private readonly sql: Sql) {}

  static async create (ctx: MeasureContext, sql: Sql): Promise<PostgresDB> {
    const db = new PostgresDB(sql)
    await db.initSchema(ctx)
    return db
  }

  async execute<T extends any[] = (Row & Iterable<Row>)[]>(query: string, params?: any[]): Promise<T> {
    query = params !== undefined && params.length > 0 ? injectVars(query, params) : query
    return await this.sql.unsafe<T>(query)
  }

  async initSchema (ctx: MeasureContext): Promise<void> {
    await this.execute('CREATE SCHEMA IF NOT EXISTS billing')
    await this.execute(`
        CREATE TABLE IF NOT EXISTS billing.migrations (
          name       VARCHAR(255) NOT NULL,
          created_on TIMESTAMP    NOT NULL DEFAULT now()
        )
    `)

    const appliedMigrations = (await this.execute<Row[]>('SELECT name FROM billing.migrations')).map((row) => row.name)
    ctx.info('applied migrations', { migrations: appliedMigrations })

    for (const [name, sql] of getMigrations()) {
      if (appliedMigrations.includes(name)) {
        continue
      }

      try {
        ctx.warn('applying migration', { migration: name })
        await this.execute(sql)
        await this.execute('INSERT INTO billing.migrations (name) VALUES ($1)', [name])
      } catch (err: any) {
        ctx.error('failed to apply migration', { migration: name, error: err })
        throw err
      }
    }
  }

  async getLiveKitStats (
    ctx: MeasureContext,
    workspace: WorkspaceUuid,
    start: Date,
    end: Date
  ): Promise<LiveKitUsageData> {
    return {
      sessions: await this.getDailySessionTotals(ctx, workspace, start, end),
      egress: await this.getDailyEgressTotals(ctx, workspace, start, end)
    }
  }

  async getDailySessionTotals (
    ctx: MeasureContext,
    workspace: WorkspaceUuid,
    start: Date,
    end: Date
  ): Promise<LiveKitSessionsUsageData[]> {
    const query = `
      SELECT
          DATE_TRUNC('day', session_start) AS day,
          COALESCE(SUM(bandwidth), 0) AS bandwidth,
          COALESCE(SUM(minutes), 0) AS minutes
      FROM billing.livekit_session
      WHERE 
          workspace = $1
          AND session_start >= $2
          AND session_start <= $3
      GROUP BY DATE_TRUNC('day', session_start)
      ORDER BY day;
    `

    const params = [workspace, start, end]

    const sessionTotals = await this.execute<{ day: string, bandwidth: string, minutes: string }[]>(query, params)
    return sessionTotals.map((s) => {
      return { day: s.day, bandwidth: parseInt(s.bandwidth), minutes: parseInt(s.minutes) }
    })
  }

  async getDailyEgressTotals (
    ctx: MeasureContext,
    workspace: WorkspaceUuid,
    start: Date,
    end: Date
  ): Promise<LiveKitEgressUsageData[]> {
    const query = `
      SELECT
          DATE_TRUNC('day', egress_start) AS day,
          COALESCE(SUM(duration), 0) AS minutes
      FROM billing.livekit_egress
      WHERE 
          workspace = $1
          AND egress_start >= $2
          AND egress_start <= $3
      GROUP BY DATE_TRUNC('day', egress_start)
      ORDER BY day;
    `

    const params = [workspace, start, end]

    const egressTotals = await this.execute<{ day: string, minutes: string }[]>(query, params)

    return egressTotals.map((e) => {
      return { day: e.day, minutes: parseInt(e.minutes) }
    })
  }

  async listLiveKitSessions (ctx: MeasureContext, workspace: WorkspaceUuid): Promise<LiveKitSessionData[] | null> {
    const query = `
        SELECT workspace, session_id, session_start, session_end, room, bandwidth, minutes
        FROM billing.livekit_session
        WHERE workspace = $1
      `
    const params = [workspace]

    return await this.execute<LiveKitSessionData[]>(query, params)
  }

  async listLiveKitEgress (ctx: MeasureContext, workspace: WorkspaceUuid): Promise<LiveKitEgressData[] | null> {
    const query = `
        SELECT workspace, egress_id, egress_start, egress_end, room, duration
        FROM billing.livekit_egress
        WHERE workspace = $1
      `
    const params = [workspace]
    return await this.execute<LiveKitEgressData[]>(query, params)
  }

  async setLiveKitSessions (ctx: MeasureContext, data: LiveKitSessionData[]): Promise<void> {
    const uniqueSessions = new Map<string, LiveKitSessionData>()
    for (const item of data) {
      uniqueSessions.set(`${item.workspace}::${item.sessionId}`, item)
    }
    const uniqueSessionsValues = uniqueSessions.values()
    for (let i = 0; i < uniqueSessions.size; i += BATCH_SIZE) {
      const batch = uniqueSessionsValues.take(BATCH_SIZE)
      const values = []
      const params = []
      let paramIndex = 1

      for (const item of batch) {
        const { workspace, sessionId, sessionStart, sessionEnd, room, bandwidth, minutes } = item
        values.push(
          `($${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++})`
        )
        params.push(workspace, sessionId, sessionStart, sessionEnd, room, bandwidth ?? 0, minutes)
      }

      if (values.length === 0) continue

      const query = `
        UPSERT INTO billing.livekit_session (workspace, session_id, session_start, session_end, room, bandwidth, minutes)
        VALUES ${values.join(',')}
      `
      await this.execute(query, params)
    }
  }

  async setLiveKitEgress (ctx: MeasureContext, data: LiveKitEgressData[]): Promise<void> {
    const uniqueSessions = new Map<string, LiveKitEgressData>()
    for (const item of data) {
      uniqueSessions.set(`${item.workspace}::${item.egressId}`, item)
    }
    const uniqueSessionsValues = uniqueSessions.values()
    for (let i = 0; i < uniqueSessions.size; i += BATCH_SIZE) {
      const batch = uniqueSessionsValues.take(BATCH_SIZE)
      const values = []
      const params = []
      let paramIndex = 1

      for (const item of batch) {
        const { workspace, egressId, egressStart, egressEnd, room, duration } = item
        values.push(
          `($${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++})`
        )
        params.push(workspace, egressId, egressStart, egressEnd, room, duration)
      }

      if (values.length === 0) continue

      const query = `
        UPSERT INTO billing.livekit_egress (workspace, egress_id, egress_start, egress_end, room, duration)
        VALUES ${values.join(',')}
      `
      await this.execute(query, params)
    }
  }

  async pushAiTranscriptData (ctx: MeasureContext, data: AiTranscriptData[]): Promise<void> {
    for (let i = 0; i < data.length; i += BATCH_SIZE) {
      const batch = data.slice(i, i + BATCH_SIZE)
      if (batch.length === 0) continue

      const values: string[] = []
      const params: any[] = []
      let paramIndex = 1

      for (const item of batch) {
        const { workspace, lastRequestId, lastStartTime, durationSeconds, usd, day } = item
        values.push(
          `($${paramIndex++}::uuid, DATE($${paramIndex++}::timestamp), $${paramIndex++}::string, $${paramIndex++}::timestamp, $${paramIndex++}::float, $${paramIndex++}::decimal)`
        )
        params.push(workspace, day, lastRequestId, lastStartTime, durationSeconds, usd)
      }

      const query = `
      INSERT INTO billing.ai_transcript_usage
        (workspace, day, last_request_id, last_start_time, total_duration_seconds, total_usd)
      VALUES ${values.join(',')}
      ON CONFLICT (workspace, day)
      DO UPDATE SET
        total_duration_seconds = billing.ai_transcript_usage.total_duration_seconds + EXCLUDED.total_duration_seconds,
        total_usd = billing.ai_transcript_usage.total_usd + EXCLUDED.total_usd,
        last_request_id = CASE
          WHEN EXCLUDED.last_start_time > billing.ai_transcript_usage.last_start_time
          THEN EXCLUDED.last_request_id
          ELSE billing.ai_transcript_usage.last_request_id
        END,
        last_start_time = GREATEST(billing.ai_transcript_usage.last_start_time, EXCLUDED.last_start_time);
    `

      await this.execute(query, params)
    }
  }

  async getAiTranscriptStats (
    ctx: MeasureContext,
    workspace: WorkspaceUuid,
    start?: Date,
    end?: Date
  ): Promise<AiTranscriptUsage> {
    const baseSql = `
    SELECT
      SUM(total_duration_seconds) AS total_duration_seconds
    FROM billing.ai_transcript_usage
  `

    let where = 'WHERE workspace = $1::uuid'
    const params: any[] = [workspace]
    let paramIndex = params.length + 1

    if (start != null) {
      const s = new Date(start)
      s.setHours(0, 0, 0, 0)

      where += ` AND day >= $${paramIndex++}::date`
      params.push(s)
    }

    if (end != null) {
      const e = new Date(end)
      e.setHours(23, 59, 59, 999)

      where += ` AND day <= $${paramIndex++}::date`
      params.push(e)
    }

    const sql = [baseSql, where].join(' ')
    const result = await this.execute(sql, params)

    return {
      totalDurationSeconds: Number(result[0]?.total_duration_seconds ?? 0)
    }
  }

  async getAiTranscriptLastData (ctx: MeasureContext): Promise<AiTranscriptData | undefined> {
    const sql = `
    SELECT *
    FROM billing.ai_transcript_usage
    ORDER BY day DESC
    LIMIT 1`

    const result = await this.execute(sql)
    const last = result[0]

    if (last == null) return undefined

    return {
      workspace: last.workspace,
      day: last.day,
      lastRequestId: last.last_request_id,
      lastStartTime: last.last_start_time,
      durationSeconds: Number(last?.total_duration_seconds ?? 0),
      usd: Number(last.usd)
    }
  }

  async pushAiTokensData (ctx: MeasureContext, data: AiTokensData[]): Promise<void> {
    const BATCH_SIZE = 100

    for (let i = 0; i < data.length; i += BATCH_SIZE) {
      const batch = data.slice(i, i + BATCH_SIZE)
      if (batch.length === 0) continue

      const values: string[] = []
      const params: any[] = []
      let paramIndex = 1

      for (const item of batch) {
        const { workspace, reason, tokens, date } = item

        values.push(
          `($${paramIndex++}::uuid, $${paramIndex++}::date, $${paramIndex++}::string, $${paramIndex++}::int8)`
        )

        params.push(workspace, date, reason, tokens)
      }

      const sql = `
      INSERT INTO billing.ai_tokens_usage (workspace, day, reason, total_tokens)
      VALUES ${values.join(',')}
      ON CONFLICT (workspace, day, reason)
      DO UPDATE SET
        total_tokens = billing.ai_tokens_usage.total_tokens + EXCLUDED.total_tokens;
    `

      await this.execute(sql, params)
    }
  }

  async getAiTokensStats (
    ctx: MeasureContext,
    workspace: WorkspaceUuid,
    start?: Date,
    end?: Date
  ): Promise<AiTokensUsage[]> {
    const baseSql = `
    SELECT
      reason,
      SUM(total_tokens) AS total_tokens
    FROM billing.ai_tokens_usage
  `

    let where = 'WHERE workspace = $1::uuid'
    const params: any[] = [workspace]
    let paramIndex = 2

    if (start != null) {
      const s = new Date(start)
      s.setHours(0, 0, 0, 0)

      where += ` AND day >= $${paramIndex++}::date`
      params.push(s)
    }

    if (end != null) {
      const e = new Date(end)
      e.setHours(23, 59, 59, 999)

      where += ` AND day <= $${paramIndex++}::date`
      params.push(e)
    }

    const groupBy = 'GROUP BY reason'
    const orderBy = 'ORDER BY reason ASC'

    const sql = [baseSql, where, groupBy, orderBy].join(' ')
    const result = await this.execute(sql, params)

    return result.map((row: any) => ({
      reason: row.reason,
      totalTokens: Number(row.total_tokens ?? 0)
    }))
  }
}

export default PostgresDB

function injectVars (sql: string, values: any[]): string {
  return sql.replaceAll(/(\$\d+)/g, (_, idx) => {
    return escape(values[parseInt(idx.substring(1)) - 1])
  })
}

function escape (value: any): string {
  if (value === null || value === undefined) {
    return 'NULL'
  }

  if (Array.isArray(value)) {
    return 'ARRAY[' + value.map(escape).join(',') + ']'
  }

  switch (typeof value) {
    case 'number':
      if (isNaN(value) || !isFinite(value)) {
        throw new Error('Invalid number value')
      }
      return value.toString()
    case 'boolean':
      return value ? 'TRUE' : 'FALSE'
    case 'string':
      return `'${value.replace(/'/g, "''")}'`
    case 'object':
      if (value instanceof Date) {
        return `'${value.toISOString()}'`
      } else {
        return `'${JSON.stringify(value)}'`
      }
    default:
      throw new Error(`Unsupported value type: ${typeof value}`)
  }
}
